package com.xuan.judge.config;


import cn.hutool.core.io.FileUtil;
import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.command.*;
import com.github.dockerjava.api.model.*;
import com.xuan.common.core.constants.JudgeConstants;
import org.springframework.util.CollectionUtils;
import java.io.File;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;



public class DockerSandBoxPool {


    private final DockerClient dockerClient;

    private final Long memoryLimit;

    private final Long memorySwapLimit;

    private final Long cpuLimit;

    private final String jdkImage;

    private final String containerNamePrefix;

    private final int poolSize;

    // 存储 containerId
    private final BlockingQueue<String> containerIdQueue;
                                    // 容器id-容器名称
    private final Map<String,String> containerNameMap;


    public DockerSandBoxPool(DockerClient dockerClient, Long memoryLimit,
                             Long memorySwapLimit, Long cpuLimit,
                             String jdkImage, String containerNamePrefix, int poolSize) {
        this.dockerClient = dockerClient;
        this.memoryLimit = memoryLimit;
        this.memorySwapLimit = memorySwapLimit;
        this.cpuLimit = cpuLimit;
        this.jdkImage = jdkImage;
        this.containerNamePrefix = containerNamePrefix;
        this.poolSize = poolSize;
        containerIdQueue = new ArrayBlockingQueue<>(poolSize);
        containerNameMap = new ConcurrentHashMap<>();
        initDockerPool();
    }


    private void initDockerPool() {
        // 使用 dockerClient 拉取 docker 镜像(如果不存在)
        pullJavaEnvImage();
        for (int i = 0; i < poolSize; i++) {
            // 前缀 + 序号
            createContainer(containerNamePrefix + "-" + i);
        }

    }

    public String getContainer() {
        try {
            return containerIdQueue.take();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    public void returnContainer(String containerId) {
        try {
            containerIdQueue.put(containerId);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }


    public String getContainerDir(String containerId) {
        String containerName = containerNameMap.get(containerId);
        return System.getProperty("user.dir") + File.separator
                + JudgeConstants.QUESTION_CODE_DIR_POOL + File.separator
                + containerName;
    }


    private void createContainer(String containerName) {
        // 1. 判断容器是否已存在, 没有再创建
        List<Container> containerList = dockerClient.listContainersCmd().withShowAll(true).exec();// 获取所有已启动和未启动的容器
        if (!CollectionUtils.isEmpty(containerList)) {
            // docker 获取的容器名称带有前缀 /
            String name = JudgeConstants.DOCKER_CONTAINS_PREFIX + containerName;
            for (Container container : containerList) {
                // 容器已存在
                if (Objects.equals(container.getNames()[0], name)) {
                    if ("created".equalsIgnoreCase(container.getState()) || "exited".equalsIgnoreCase(container.getState())) {
                        // 启动容器
                        dockerClient.startContainerCmd(container.getId()).exec();
                    }
                    containerIdQueue.add(container.getId());
                    containerNameMap.put(container.getId(), containerName);
                    return;
                }
            }
        }


        // 指定创建容器的镜像和 容器名称
        CreateContainerCmd createContainerCmd = dockerClient.createContainerCmd(jdkImage)
                .withName(containerName);
        // 指定容器启动参数: 限制资源, 控制权限
        HostConfig hostConfig = getHostConfig(containerName);
        // 创建容器
        CreateContainerResponse containerResponse = createContainerCmd.withHostConfig(hostConfig)
                .withAttachStdout(true)         // 配置标准输出:
                .withAttachStderr(true)         // 标准错误输出:在命令执行出错时获取命令行的错误提示
                .withTty(true)                  // 为容器配置一个虚终端
                .exec();
        // 记录容器 ID
        String containerId = containerResponse.getId();
        containerIdQueue.add(containerId);
        containerNameMap.put(containerId, containerName);
        // 启动容器
        dockerClient.startContainerCmd(containerId).exec();
    }


    /**
     *  拉取 docker 镜像,若不存在
     */
    private void pullJavaEnvImage() {
        // 1. 先判断要拉取的镜像是否存在
        ListImagesCmd listImagesCmd = dockerClient.listImagesCmd();
        List<Image> imageList = listImagesCmd.exec();
        for (Image image : imageList) {
            String[] repoTags = image.getRepoTags();
            if (repoTags != null && repoTags.length > 0 && jdkImage.equals(repoTags[0])) {
                return;
            }
        }
        // 2. 不存在则拉取
        PullImageCmd pullImageCmd = dockerClient.pullImageCmd(JudgeConstants.JAVA_ENV_IMAGE);
        try {
            pullImageCmd.exec(new PullImageResultCallback()).awaitCompletion();     // 阻塞等待镜像拉取完成
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }


    /**
     *  对容器相关参数的配置
     *  比如: 对资源,网络访问和文件读写权限的限制
     * @return
     */
    private HostConfig getHostConfig(String containerName) {
        HostConfig hostConfig = new HostConfig();
        // 指的文件的挂载目录: 每个容器对应一个挂载目录(按容器名称)
        String containerDir = createContainerDir(containerName);
        hostConfig.setBinds(new Bind(containerDir, new Volume(JudgeConstants.DOCKER_USER_CODE_DIR)));
        hostConfig
                .withMemory(memoryLimit)            // 分配给容器的内存
                .withMemorySwap(memorySwapLimit)    // 内存交换区
                .withCpuCount(cpuLimit)             // CPU 数量
                .withNetworkMode("none")            // 禁用网络
                .withReadonlyRootfs(true);          // 限制文件只读权限
        return hostConfig;
    }


    private String createContainerDir(String containerName) {
        // 1. 先创建一级目录，存在各个容器的提交的代码
        String codeDir = System.getProperty("user.dir") + File.separator + JudgeConstants.QUESTION_CODE_DIR_POOL;
        if (!FileUtil.exist(codeDir)) {
            FileUtil.mkdir(codeDir);
        }
        // 2. 为每个容器创建单独的二级目录，存放该容器下用户提交的代码
        return codeDir + File.separator + containerName;
    }
}
